#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/3/6 16:12
# @Author  : wanghaoran
# @File    : keywords_retrieval.py
from base import BaseRetrievalEngineTool
from elasticsearch_dsl import Search
import pandas as pd


class KeywordsRetrieval(BaseRetrievalEngineTool):
    def __init__(self):
        pass

    def get_results(self, query, base_name, topk):
        s = Search().query('multi_match', query=query, fields=['text', 'detail_text'])
        s = s.query('terms', base__name=[base_name])

        es_res = s.execute()
        es_results = [hit for hit in es_res.hits.hits]
        keyword_results = []
        for r in es_results[:topk]:
            keyword_results.append({
                "uid": r["_source"]["uid"],
                "question": r["_source"]["text"],
                "detail_text": r["_source"]["detail_text"],
                "score": r["_score"],
                "comment": "keyword_es"
            })
        keyword_results_df = pd.DataFrame(keyword_results)
        return keyword_results_df
